# $Id: stdlib.py e501b84d8043 2013/07/05 07:47:14 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#
__version__  = "$Revision: e501b84d8043 $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2013-07-05 07:47:14Z $"

# essential imports
import array
import random
import socket
import string
import struct
import sys
import StringIO

IPv4MaxValue = 32
BgpAsMax = 65536

COLORS = {
      'grey'    : '\033[90m',
      'red'     : '\033[91m',
      'green'   : '\033[92m',
      'yellow'  : '\033[93m',
      'blue'    : '\033[94m',
      'magenta' : '\033[95m',
      'cyan'    : '\033[96m',
      }

my_str = lambda value: str(value).encode('utf-8')
my_class = lambda klass: klass.__class__
my_name = lambda klass: my_str(my_class(klass).__name__)
packed2ip = lambda bits: socket.inet_ntoa(struct.pack(">I", bits))
ip2packed = lambda ip: struct.unpack(">I", socket.inet_aton(ip))[0]
randstring = lambda value: GetRandomString(value, False)

def fprintf(output, fmt, *args):
   output.write(fmt %args)

def printf(fmt, *args):
   fprintf(sys.stdout, fmt, *args)

def error(fmt, *args):
   printf(fmt, *args)
   sys.exit(1)

def sprintf(fmt, *args):
   newstr = StringIO.StringIO()
   newstr.write(fmt %args)
   buffer = newstr.getvalue()
   newstr.close()

   return buffer

def netmask4(Number, Wildcard = False):
   retval = None
   if isinstance(Number, int):
      if Number > IPv4MaxValue:
         raise ValueError, "too big a value to get subnet information"
      elif Number < 0:
         raise ValueError, "subnet value must be postive"
      else:
         if Wildcard:
            bits = 0x00000000 ^ (1 << IPv4MaxValue - Number) - 1
         else:
            bits = 0xffffffff ^ (1 << IPv4MaxValue - Number) - 1

         retval = packed2ip(bits)
   else:
      raise ValueError, "integer value expected, got %s" %type(Number)

   return unicode(retval)

## Function that returns random string based on the length provided
def GetRandomString(Length, unicode = True):
   """
   Return random string based on the length provided in the argument
   AssertionError is raised if the generated string length is not equal to
   the asked string length
   @param Length: Length of the random string needed
   @return random string
   """
   retval = None
   if Length:
      chars = string.printable[0:62]
      length = 0
      try:
         length = int(Length)
      except ValueError, message:
         raise ValueError, message
      else:
         retval = ''.join([random.choice(chars) for ix in range(length)])

   assert len(retval) == length, "string length %d shorter than %d" %(
         len(retval), Length)

   if unicode:
      retval = unicode(retval.encode('utf-8'))

   return retval

def pack2ip(value):
   maxval = 2 << 31
   if value > maxval:
      raise ValueError, "Cannot unpack value greater thatn 2 ^ 31"

   octets = []
   for _ in xrange(4):
      octets.insert(0, str(value & 0xFF))
      value >>= 8

   return '.'.join(octets)

def ip2pack(address):
   retval = 0
   count = 0
   for ix in address.split("."):
      val = int(ix)
      if count == 0:
         retval |= val << 24
      elif count == 1:
         retval |= val << 16
      elif count == 2:
         retval |= val << 8
      elif count == 3:
         retval |= val

      count += 1

   return int(retval)

def _getprime(number):
   retval = None
   if number < 0:
      raise ValueError, "Invalid number"
   elif number == 1:
      retval = number
   elif number == 2:
      retval = number
   else:
      retval = []
      nums = list(xrange(3, (number + 1), 2))
      nlen = (number // 2) - 1 + (number % 2)
      idx = 0
      idxsqrt = (int(number ** 0.5) - 3) // 2
      while idx <= idxsqrt:
         nidx = (idx << 1) + 3
         value = xrange(idx * (nidx + 3) + 3, nlen, nidx)
         for jy in value:
            nums[jy] = 0

         idx += 1
         while idx <= idxsqrt:
            if nums[idx] != 0:
               break

            idx += 1

      retval = [2] + [ix for ix in nums if ix != 0]

   return retval

def getprime(number, start = 2):
   retval = []
   if number < start:
      raise TypeError("starting number %d cannot be higher than %d" %(
            start, number))

   primelist = _getprime(number)
   if isinstance(primelist, list):
      for ix in primelist:
         if ix >= start:
            retval.append(ix)
   else:
      retval.append(primelist)

   return retval

def isprime(number):
   retval = None
   number = abs(int(number))
   if number < 2:
      retval = False
   elif number == 2:
      retval = True
   elif not number & 1:
      retval = False
   else:
      value = xrange(3, int(number ** 0.5) + 1, 2)
      for ix in value:
         if (number % ix) == 0:
            retval = False
            break

      if retval is None:
         retval = True

   return retval

def ip2ptr(address):
   iplist = str(address).split(".")
   iplist.reverse()
   return ".".join(iplist)

def ip2ptr6(address):
   retval = None
   try:
      addr = array.array('L', struct.unpack('>4L',
         socket.inet_pton(socket.AF_INET6, address)))
      addr.reverse()
      address = sum(addr[i] << (i * 32) for i in range(len(addr)))
   except socket.error:
      pass
   except struct.error:
      pass
   else:
      temp = "%32x" %int(address)
      iplist = [ix for ix in temp]
      iplist.reverse()
      retval = ".".join(iplist)

   return retval

class EnumType(object):
   def __init__(self, *args, **kwargs):
      self._enum = self.enum(*args, **kwargs)
      self.__temp = self._enum.__dict__
      self._count = len(args) + len(kwargs)

   def enum(self, *sequential, **named):
      enums = dict(zip(sequential, xrange(len(sequential))), **named)
      return type('Enum', (), enums)

   def __getattr__(self, name):
      retval = None
      if isinstance(name, (str, unicode)):
         if name in self:
            retval = self.__temp[name]

      return retval

   def __contains__(self, name):
      retval = False
      if name in self.__temp.keys():
         retval = True

      return retval

   def __len__(self):
      return self._count

   def __getitem__(self, Key):
      retval = None
      if isinstance(Key, (int, long)):
         for (key, value) in self.__temp.items():
            if isinstance(value, (int, long)) and value == Key:
               retval = key
               break
         
         retval = retval.encode('utf-8')
      elif isinstance(Key, (str, unicode)):
         retval = getattr(self, Key)

      return retval

def num2asdot(number):
   retval = ""
   if number:
      retval = str(number / BgpAsMax)
      if int(retval) == 0:
         retval = str(number % BgpAsMax)
      else:
         retval += "." + str(number % BgpAsMax)

   if not retval:
      retval = 0

   return retval

def asdot2num(asnumber):
   asnumber = str(asnumber)
   retval = 0
   length = len(asnumber.split("."))
   if length == 2:
      (val, val1) = asnumber.split(".")
      retval = (int(val) * BgpAsMax) + int(val1)
   elif length == 1:
      retval = int(asnumber)

   return retval

def colorprint(color, fmt, *args):
   errmesg = None
   if len(args):
      errmesg = sprintf(fmt, *args)
   else:
      errmesg = str(fmt)

   if color in COLORS:
      endc = '\033[0m'
      print(COLORS[color] + errmesg + endc)
   else:
      print(errmesg)
